import mindspore as ms
import argparse
import numpy as np
import torch
import sys
sys.path.append("/home/data/zhangyh22/ABINet/ABINet_Mindspore/abinet_gpu/abinet_ascend")
from mindspore import save_checkpoint, Tensor
#from convert_utils import print_dict, generate_params_dict
from print_para import mindspore_params,pytorch_params_convert_weight
from ms_utils import Config
from src.ms_model_abinet_iter import ABINetIterModel as ms_ABINetIterModel
# from modules.model_abinet_iter import ABINetIterModel
from src.ms_BaseVision import BaseVision as ms_BaseVision
from src.ms_resnet import resnet45 as ms_resnet50

def split_torch_attention(state):
    s = list(state.keys())
    for name in s:
        if name.endswith('self_attn.in_proj_weight') or name.endswith('self_attn.in_proj_bias') or name.endswith('multihead_attn.in_proj_weight')  or name.endswith('multihead_attn.in_proj_bias'):
            value = state.pop(name)
            q, k, v = np.split(value, 3, 0)
            state[name + '.q'] = ms.Tensor(q)
            state[name + '.k'] = ms.Tensor(k)
            state[name + '.v'] = ms.Tensor(v)


bn_ms2pt = {"gamma": "weight",
            "beta": "bias",
            "moving_mean": "running_mean",
            "moving_variance": "running_var"}

def param_convert(ms_params, pt_params, ckpt_path):
    # 参数名映射字典
    param = {
    'vision.backbone.resnet.conv1.weight':'vision.backbone.resnet.conv1.weight',
    'vision.backbone.resnet.bn1.moving_mean' :'vision.backbone.resnet.bn1.running_mean',
    'vision.backbone.resnet.bn1.moving_variance':'vision.backbone.resnet.bn1.running_var', 
    'vision.backbone.resnet.bn1.gamma':'vision.backbone.resnet.bn1.weight',
    'vision.backbone.resnet.bn1.beta':'vision.backbone.resnet.bn1.bias',
    'vision.backbone.resnet.layer1.0.conv1.weight':'vision.backbone.resnet.layer1.0.conv1.weight',
    'vision.backbone.resnet.layer1.0.bn1.moving_mean':'vision.backbone.resnet.layer1.0.bn1.running_mean',
    'vision.backbone.resnet.layer1.0.bn1.moving_variance':'vision.backbone.resnet.layer1.0.bn1.running_var',
    'vision.backbone.resnet.layer1.0.bn1.gamma':'vision.backbone.resnet.layer1.0.bn1.weight',
    'vision.backbone.resnet.layer1.0.bn1.beta':'vision.backbone.resnet.layer1.0.bn1.bias',
    'vision.backbone.resnet.layer1.0.conv2.weight':'vision.backbone.resnet.layer1.0.conv2.weight',
    'vision.backbone.resnet.layer1.0.bn2.moving_mean':'vision.backbone.resnet.layer1.0.bn2.running_mean',
    'vision.backbone.resnet.layer1.0.bn2.moving_variance':'vision.backbone.resnet.layer1.0.bn2.running_var',
    'vision.backbone.resnet.layer1.0.bn2.gamma':'vision.backbone.resnet.layer1.0.bn2.weight',
    'vision.backbone.resnet.layer1.0.bn2.beta':'vision.backbone.resnet.layer1.0.bn2.bias',
    'vision.backbone.resnet.layer1.0.downsample.0.weight':'vision.backbone.resnet.layer1.0.downsample.0.weight',
    'vision.backbone.resnet.layer1.0.downsample.1.moving_mean':'vision.backbone.resnet.layer1.0.downsample.1.running_mean',
    'vision.backbone.resnet.layer1.0.downsample.1.moving_variance':'vision.backbone.resnet.layer1.0.downsample.1.running_var',
    'vision.backbone.resnet.layer1.0.downsample.1.gamma':'vision.backbone.resnet.layer1.0.downsample.1.weight',
    'vision.backbone.resnet.layer1.0.downsample.1.beta':'vision.backbone.resnet.layer1.0.downsample.1.bias',
    'vision.backbone.resnet.layer1.1.conv1.weight':'vision.backbone.resnet.layer1.1.conv1.weight',
    'vision.backbone.resnet.layer1.1.bn1.moving_mean':'vision.backbone.resnet.layer1.1.bn1.running_mean',
    'vision.backbone.resnet.layer1.1.bn1.moving_variance':'vision.backbone.resnet.layer1.1.bn1.running_var',
    'vision.backbone.resnet.layer1.1.bn1.gamma':'vision.backbone.resnet.layer1.1.bn1.weight',
    'vision.backbone.resnet.layer1.1.bn1.beta':'vision.backbone.resnet.layer1.1.bn1.bias',
    'vision.backbone.resnet.layer1.1.conv2.weight':'vision.backbone.resnet.layer1.1.conv2.weight',
    'vision.backbone.resnet.layer1.1.bn2.moving_mean':'vision.backbone.resnet.layer1.1.bn2.running_mean',
    'vision.backbone.resnet.layer1.1.bn2.moving_variance':'vision.backbone.resnet.layer1.1.bn2.running_var',
    'vision.backbone.resnet.layer1.1.bn2.gamma':'vision.backbone.resnet.layer1.1.bn2.weight',
    'vision.backbone.resnet.layer1.1.bn2.beta':'vision.backbone.resnet.layer1.1.bn2.bias',
    'vision.backbone.resnet.layer1.2.conv1.weight':'vision.backbone.resnet.layer1.2.conv1.weight',
    'vision.backbone.resnet.layer1.2.bn1.moving_mean':'vision.backbone.resnet.layer1.2.bn1.running_mean',
    'vision.backbone.resnet.layer1.2.bn1.moving_variance':'vision.backbone.resnet.layer1.2.bn1.running_var',
    'vision.backbone.resnet.layer1.2.bn1.gamma':'vision.backbone.resnet.layer1.2.bn1.weight',
    'vision.backbone.resnet.layer1.2.bn1.beta':'vision.backbone.resnet.layer1.2.bn1.bias',
    'vision.backbone.resnet.layer1.2.conv2.weight':'vision.backbone.resnet.layer1.2.conv2.weight',
    'vision.backbone.resnet.layer1.2.bn2.moving_mean':'vision.backbone.resnet.layer1.2.bn2.running_mean',
    'vision.backbone.resnet.layer1.2.bn2.moving_variance':'vision.backbone.resnet.layer1.2.bn2.running_var',
    'vision.backbone.resnet.layer1.2.bn2.gamma':'vision.backbone.resnet.layer1.2.bn2.weight',
    'vision.backbone.resnet.layer1.2.bn2.beta':'vision.backbone.resnet.layer1.2.bn2.bias',
    'vision.backbone.resnet.layer2.0.conv1.weight':'vision.backbone.resnet.layer2.0.conv1.weight',
    'vision.backbone.resnet.layer2.0.bn1.moving_mean':'vision.backbone.resnet.layer2.0.bn1.running_mean',
    'vision.backbone.resnet.layer2.0.bn1.moving_variance':'vision.backbone.resnet.layer2.0.bn1.running_var',
    'vision.backbone.resnet.layer2.0.bn1.gamma':'vision.backbone.resnet.layer2.0.bn1.weight',
    'vision.backbone.resnet.layer2.0.bn1.beta':'vision.backbone.resnet.layer2.0.bn1.bias',
    'vision.backbone.resnet.layer2.0.conv2.weight':'vision.backbone.resnet.layer2.0.conv2.weight',
    'vision.backbone.resnet.layer2.0.bn2.moving_mean':'vision.backbone.resnet.layer2.0.bn2.running_mean',
    'vision.backbone.resnet.layer2.0.bn2.moving_variance':'vision.backbone.resnet.layer2.0.bn2.running_var',
    'vision.backbone.resnet.layer2.0.bn2.gamma':'vision.backbone.resnet.layer2.0.bn2.weight',
    'vision.backbone.resnet.layer2.0.bn2.beta':'vision.backbone.resnet.layer2.0.bn2.bias',
    'vision.backbone.resnet.layer2.0.downsample.0.weight':'vision.backbone.resnet.layer2.0.downsample.0.weight',
    'vision.backbone.resnet.layer2.0.downsample.1.moving_mean':'vision.backbone.resnet.layer2.0.downsample.1.running_mean',
    'vision.backbone.resnet.layer2.0.downsample.1.moving_variance':'vision.backbone.resnet.layer2.0.downsample.1.running_var',
    'vision.backbone.resnet.layer2.0.downsample.1.gamma':'vision.backbone.resnet.layer2.0.downsample.1.weight',
    'vision.backbone.resnet.layer2.0.downsample.1.beta':'vision.backbone.resnet.layer2.0.downsample.1.bias',
    'vision.backbone.resnet.layer2.1.conv1.weight':'vision.backbone.resnet.layer2.1.conv1.weight',
    'vision.backbone.resnet.layer2.1.bn1.moving_mean':'vision.backbone.resnet.layer2.1.bn1.running_mean',
    'vision.backbone.resnet.layer2.1.bn1.moving_variance':'vision.backbone.resnet.layer2.1.bn1.running_var',
    'vision.backbone.resnet.layer2.1.bn1.gamma':'vision.backbone.resnet.layer2.1.bn1.weight',
    'vision.backbone.resnet.layer2.1.bn1.beta':'vision.backbone.resnet.layer2.1.bn1.bias',
    'vision.backbone.resnet.layer2.1.conv2.weight':'vision.backbone.resnet.layer2.1.conv2.weight',
    'vision.backbone.resnet.layer2.1.bn2.moving_mean':'vision.backbone.resnet.layer2.1.bn2.running_mean',
    'vision.backbone.resnet.layer2.1.bn2.moving_variance':'vision.backbone.resnet.layer2.1.bn2.running_var',
    'vision.backbone.resnet.layer2.1.bn2.gamma':'vision.backbone.resnet.layer2.1.bn2.weight',
    'vision.backbone.resnet.layer2.1.bn2.beta':'vision.backbone.resnet.layer2.1.bn2.bias',
    'vision.backbone.resnet.layer2.2.conv1.weight':'vision.backbone.resnet.layer2.2.conv1.weight',
    'vision.backbone.resnet.layer2.2.bn1.moving_mean':'vision.backbone.resnet.layer2.2.bn1.running_mean',
    'vision.backbone.resnet.layer2.2.bn1.moving_variance':'vision.backbone.resnet.layer2.2.bn1.running_var',
    'vision.backbone.resnet.layer2.2.bn1.gamma':'vision.backbone.resnet.layer2.2.bn1.weight',
    'vision.backbone.resnet.layer2.2.bn1.beta':'vision.backbone.resnet.layer2.2.bn1.bias',
    'vision.backbone.resnet.layer2.2.conv2.weight':'vision.backbone.resnet.layer2.2.conv2.weight',
    'vision.backbone.resnet.layer2.2.bn2.moving_mean':'vision.backbone.resnet.layer2.2.bn2.running_mean',
    'vision.backbone.resnet.layer2.2.bn2.moving_variance':'vision.backbone.resnet.layer2.2.bn2.running_var',
    'vision.backbone.resnet.layer2.2.bn2.gamma':'vision.backbone.resnet.layer2.2.bn2.weight',
    'vision.backbone.resnet.layer2.2.bn2.beta':'vision.backbone.resnet.layer2.2.bn2.bias',
    'vision.backbone.resnet.layer2.3.conv1.weight':'vision.backbone.resnet.layer2.3.conv1.weight',
    'vision.backbone.resnet.layer2.3.bn1.moving_mean':'vision.backbone.resnet.layer2.3.bn1.running_mean',
    'vision.backbone.resnet.layer2.3.bn1.moving_variance':'vision.backbone.resnet.layer2.3.bn1.running_var',
    'vision.backbone.resnet.layer2.3.bn1.gamma':'vision.backbone.resnet.layer2.3.bn1.weight',
    'vision.backbone.resnet.layer2.3.bn1.beta':'vision.backbone.resnet.layer2.3.bn1.bias',
    'vision.backbone.resnet.layer2.3.conv2.weight':'vision.backbone.resnet.layer2.3.conv2.weight',
    'vision.backbone.resnet.layer2.3.bn2.moving_mean':'vision.backbone.resnet.layer2.3.bn2.running_mean',
    'vision.backbone.resnet.layer2.3.bn2.moving_variance':'vision.backbone.resnet.layer2.3.bn2.running_var',
    'vision.backbone.resnet.layer2.3.bn2.gamma':'vision.backbone.resnet.layer2.3.bn2.weight',
    'vision.backbone.resnet.layer2.3.bn2.beta':'vision.backbone.resnet.layer2.3.bn2.bias',
    'vision.backbone.resnet.layer3.0.conv1.weight':'vision.backbone.resnet.layer3.0.conv1.weight',
    'vision.backbone.resnet.layer3.0.bn1.moving_mean':'vision.backbone.resnet.layer3.0.bn1.running_mean',
    'vision.backbone.resnet.layer3.0.bn1.moving_variance':'vision.backbone.resnet.layer3.0.bn1.running_var',
    'vision.backbone.resnet.layer3.0.bn1.gamma':'vision.backbone.resnet.layer3.0.bn1.weight',
    'vision.backbone.resnet.layer3.0.bn1.beta':'vision.backbone.resnet.layer3.0.bn1.bias',
    'vision.backbone.resnet.layer3.0.conv2.weight':'vision.backbone.resnet.layer3.0.conv2.weight',
    'vision.backbone.resnet.layer3.0.bn2.moving_mean':'vision.backbone.resnet.layer3.0.bn2.running_mean',
    'vision.backbone.resnet.layer3.0.bn2.moving_variance':'vision.backbone.resnet.layer3.0.bn2.running_var',
    'vision.backbone.resnet.layer3.0.bn2.gamma':'vision.backbone.resnet.layer3.0.bn2.weight',
    'vision.backbone.resnet.layer3.0.bn2.beta':'vision.backbone.resnet.layer3.0.bn2.bias',
    'vision.backbone.resnet.layer3.0.downsample.0.weight':'vision.backbone.resnet.layer3.0.downsample.0.weight',
    'vision.backbone.resnet.layer3.0.downsample.1.moving_mean':'vision.backbone.resnet.layer3.0.downsample.1.running_mean',
    'vision.backbone.resnet.layer3.0.downsample.1.moving_variance':'vision.backbone.resnet.layer3.0.downsample.1.running_var',
    'vision.backbone.resnet.layer3.0.downsample.1.gamma':'vision.backbone.resnet.layer3.0.downsample.1.weight',
    'vision.backbone.resnet.layer3.0.downsample.1.beta':'vision.backbone.resnet.layer3.0.downsample.1.bias',
    'vision.backbone.resnet.layer3.1.conv1.weight':'vision.backbone.resnet.layer3.1.conv1.weight',
    'vision.backbone.resnet.layer3.1.bn1.moving_mean':'vision.backbone.resnet.layer3.1.bn1.running_mean',
    'vision.backbone.resnet.layer3.1.bn1.moving_variance':'vision.backbone.resnet.layer3.1.bn1.running_var',
    'vision.backbone.resnet.layer3.1.bn1.gamma':'vision.backbone.resnet.layer3.1.bn1.weight',
    'vision.backbone.resnet.layer3.1.bn1.beta':'vision.backbone.resnet.layer3.1.bn1.bias',
    'vision.backbone.resnet.layer3.1.conv2.weight':'vision.backbone.resnet.layer3.1.conv2.weight',
    'vision.backbone.resnet.layer3.1.bn2.moving_mean':'vision.backbone.resnet.layer3.1.bn2.running_mean',
    'vision.backbone.resnet.layer3.1.bn2.moving_variance':'vision.backbone.resnet.layer3.1.bn2.running_var',
    'vision.backbone.resnet.layer3.1.bn2.gamma':'vision.backbone.resnet.layer3.1.bn2.weight',
    'vision.backbone.resnet.layer3.1.bn2.beta':'vision.backbone.resnet.layer3.1.bn2.bias',
    'vision.backbone.resnet.layer3.2.conv1.weight':'vision.backbone.resnet.layer3.2.conv1.weight',
    'vision.backbone.resnet.layer3.2.bn1.moving_mean':'vision.backbone.resnet.layer3.2.bn1.running_mean',
    'vision.backbone.resnet.layer3.2.bn1.moving_variance':'vision.backbone.resnet.layer3.2.bn1.running_var',
    'vision.backbone.resnet.layer3.2.bn1.gamma':'vision.backbone.resnet.layer3.2.bn1.weight',
    'vision.backbone.resnet.layer3.2.bn1.beta':'vision.backbone.resnet.layer3.2.bn1.bias',
    'vision.backbone.resnet.layer3.2.conv2.weight':'vision.backbone.resnet.layer3.2.conv2.weight',
    'vision.backbone.resnet.layer3.2.bn2.moving_mean':'vision.backbone.resnet.layer3.2.bn2.running_mean',
    'vision.backbone.resnet.layer3.2.bn2.moving_variance':'vision.backbone.resnet.layer3.2.bn2.running_var',
    'vision.backbone.resnet.layer3.2.bn2.gamma':'vision.backbone.resnet.layer3.2.bn2.weight',
    'vision.backbone.resnet.layer3.2.bn2.beta':'vision.backbone.resnet.layer3.2.bn2.bias',
    'vision.backbone.resnet.layer3.3.conv1.weight':'vision.backbone.resnet.layer3.3.conv1.weight',
    'vision.backbone.resnet.layer3.3.bn1.moving_mean':'vision.backbone.resnet.layer3.3.bn1.running_mean',
    'vision.backbone.resnet.layer3.3.bn1.moving_variance':'vision.backbone.resnet.layer3.3.bn1.running_var',
    'vision.backbone.resnet.layer3.3.bn1.gamma':'vision.backbone.resnet.layer3.3.bn1.weight',
    'vision.backbone.resnet.layer3.3.bn1.beta':'vision.backbone.resnet.layer3.3.bn1.bias',
    'vision.backbone.resnet.layer3.3.conv2.weight':'vision.backbone.resnet.layer3.3.conv2.weight',
    'vision.backbone.resnet.layer3.3.bn2.moving_mean':'vision.backbone.resnet.layer3.3.bn2.running_mean',
    'vision.backbone.resnet.layer3.3.bn2.moving_variance':'vision.backbone.resnet.layer3.3.bn2.running_var',
    'vision.backbone.resnet.layer3.3.bn2.gamma':'vision.backbone.resnet.layer3.3.bn2.weight',
    'vision.backbone.resnet.layer3.3.bn2.beta':'vision.backbone.resnet.layer3.3.bn2.bias',
    'vision.backbone.resnet.layer3.4.conv1.weight':'vision.backbone.resnet.layer3.4.conv1.weight',
    'vision.backbone.resnet.layer3.4.bn1.moving_mean':'vision.backbone.resnet.layer3.4.bn1.running_mean',
    'vision.backbone.resnet.layer3.4.bn1.moving_variance':'vision.backbone.resnet.layer3.4.bn1.running_var',
    'vision.backbone.resnet.layer3.4.bn1.gamma':'vision.backbone.resnet.layer3.4.bn1.weight',
    'vision.backbone.resnet.layer3.4.bn1.beta':'vision.backbone.resnet.layer3.4.bn1.bias',
    'vision.backbone.resnet.layer3.4.conv2.weight':'vision.backbone.resnet.layer3.4.conv2.weight',
    'vision.backbone.resnet.layer3.4.bn2.moving_mean':'vision.backbone.resnet.layer3.4.bn2.running_mean',
    'vision.backbone.resnet.layer3.4.bn2.moving_variance':'vision.backbone.resnet.layer3.4.bn2.running_var',
    'vision.backbone.resnet.layer3.4.bn2.gamma':'vision.backbone.resnet.layer3.4.bn2.weight',
    'vision.backbone.resnet.layer3.4.bn2.beta':'vision.backbone.resnet.layer3.4.bn2.bias',
    'vision.backbone.resnet.layer3.5.conv1.weight':'vision.backbone.resnet.layer3.5.conv1.weight',
    'vision.backbone.resnet.layer3.5.bn1.moving_mean':'vision.backbone.resnet.layer3.5.bn1.running_mean',
    'vision.backbone.resnet.layer3.5.bn1.moving_variance':'vision.backbone.resnet.layer3.5.bn1.running_var',
    'vision.backbone.resnet.layer3.5.bn1.gamma':'vision.backbone.resnet.layer3.5.bn1.weight',
    'vision.backbone.resnet.layer3.5.bn1.beta':'vision.backbone.resnet.layer3.5.bn1.bias',
    'vision.backbone.resnet.layer3.5.conv2.weight':'vision.backbone.resnet.layer3.5.conv2.weight',
    'vision.backbone.resnet.layer3.5.bn2.moving_mean':'vision.backbone.resnet.layer3.5.bn2.running_mean',
    'vision.backbone.resnet.layer3.5.bn2.moving_variance':'vision.backbone.resnet.layer3.5.bn2.running_var',
    'vision.backbone.resnet.layer3.5.bn2.gamma':'vision.backbone.resnet.layer3.5.bn2.weight',
    'vision.backbone.resnet.layer3.5.bn2.beta':'vision.backbone.resnet.layer3.5.bn2.bias',
    'vision.backbone.resnet.layer4.0.conv1.weight':'vision.backbone.resnet.layer4.0.conv1.weight',
    'vision.backbone.resnet.layer4.0.bn1.moving_mean':'vision.backbone.resnet.layer4.0.bn1.running_mean',
    'vision.backbone.resnet.layer4.0.bn1.moving_variance':'vision.backbone.resnet.layer4.0.bn1.running_var',
    'vision.backbone.resnet.layer4.0.bn1.gamma':'vision.backbone.resnet.layer4.0.bn1.weight',
    'vision.backbone.resnet.layer4.0.bn1.beta':'vision.backbone.resnet.layer4.0.bn1.bias',
    'vision.backbone.resnet.layer4.0.conv2.weight':'vision.backbone.resnet.layer4.0.conv2.weight',
    'vision.backbone.resnet.layer4.0.bn2.moving_mean':'vision.backbone.resnet.layer4.0.bn2.running_mean',
    'vision.backbone.resnet.layer4.0.bn2.moving_variance':'vision.backbone.resnet.layer4.0.bn2.running_var',
    'vision.backbone.resnet.layer4.0.bn2.gamma':'vision.backbone.resnet.layer4.0.bn2.weight',
    'vision.backbone.resnet.layer4.0.bn2.beta':'vision.backbone.resnet.layer4.0.bn2.bias',
    'vision.backbone.resnet.layer4.0.downsample.0.weight':'vision.backbone.resnet.layer4.0.downsample.0.weight',
    'vision.backbone.resnet.layer4.0.downsample.1.moving_mean':'vision.backbone.resnet.layer4.0.downsample.1.running_mean',
    'vision.backbone.resnet.layer4.0.downsample.1.moving_variance':'vision.backbone.resnet.layer4.0.downsample.1.running_var',
    'vision.backbone.resnet.layer4.0.downsample.1.gamma':'vision.backbone.resnet.layer4.0.downsample.1.weight',
    'vision.backbone.resnet.layer4.0.downsample.1.beta':'vision.backbone.resnet.layer4.0.downsample.1.bias',
    'vision.backbone.resnet.layer4.1.conv1.weight':'vision.backbone.resnet.layer4.1.conv1.weight',
    'vision.backbone.resnet.layer4.1.bn1.moving_mean':'vision.backbone.resnet.layer4.1.bn1.running_mean',
    'vision.backbone.resnet.layer4.1.bn1.moving_variance':'vision.backbone.resnet.layer4.1.bn1.running_var',
    'vision.backbone.resnet.layer4.1.bn1.gamma':'vision.backbone.resnet.layer4.1.bn1.weight',
    'vision.backbone.resnet.layer4.1.bn1.beta':'vision.backbone.resnet.layer4.1.bn1.bias',
    'vision.backbone.resnet.layer4.1.conv2.weight':'vision.backbone.resnet.layer4.1.conv2.weight',
    'vision.backbone.resnet.layer4.1.bn2.moving_mean':'vision.backbone.resnet.layer4.1.bn2.running_mean',
    'vision.backbone.resnet.layer4.1.bn2.moving_variance':'vision.backbone.resnet.layer4.1.bn2.running_var',
    'vision.backbone.resnet.layer4.1.bn2.gamma':'vision.backbone.resnet.layer4.1.bn2.weight',
    'vision.backbone.resnet.layer4.1.bn2.beta':'vision.backbone.resnet.layer4.1.bn2.bias',
    'vision.backbone.resnet.layer4.2.conv1.weight':'vision.backbone.resnet.layer4.2.conv1.weight',
    'vision.backbone.resnet.layer4.2.bn1.moving_mean':'vision.backbone.resnet.layer4.2.bn1.running_mean',
    'vision.backbone.resnet.layer4.2.bn1.moving_variance':'vision.backbone.resnet.layer4.2.bn1.running_var',
    'vision.backbone.resnet.layer4.2.bn1.gamma':'vision.backbone.resnet.layer4.2.bn1.weight',
    'vision.backbone.resnet.layer4.2.bn1.beta':'vision.backbone.resnet.layer4.2.bn1.bias',
    'vision.backbone.resnet.layer4.2.conv2.weight':'vision.backbone.resnet.layer4.2.conv2.weight',
    'vision.backbone.resnet.layer4.2.bn2.moving_mean':'vision.backbone.resnet.layer4.2.bn2.running_mean',
    'vision.backbone.resnet.layer4.2.bn2.moving_variance':'vision.backbone.resnet.layer4.2.bn2.running_var',
    'vision.backbone.resnet.layer4.2.bn2.gamma':'vision.backbone.resnet.layer4.2.bn2.weight',
    'vision.backbone.resnet.layer4.2.bn2.beta':'vision.backbone.resnet.layer4.2.bn2.bias',
    'vision.backbone.resnet.layer4.3.conv1.weight':'vision.backbone.resnet.layer4.3.conv1.weight',
    'vision.backbone.resnet.layer4.3.bn1.moving_mean':'vision.backbone.resnet.layer4.3.bn1.running_mean',
    'vision.backbone.resnet.layer4.3.bn1.moving_variance':'vision.backbone.resnet.layer4.3.bn1.running_var',
    'vision.backbone.resnet.layer4.3.bn1.gamma':'vision.backbone.resnet.layer4.3.bn1.weight',
    'vision.backbone.resnet.layer4.3.bn1.beta':'vision.backbone.resnet.layer4.3.bn1.bias',
    'vision.backbone.resnet.layer4.3.conv2.weight':'vision.backbone.resnet.layer4.3.conv2.weight',
    'vision.backbone.resnet.layer4.3.bn2.moving_mean':'vision.backbone.resnet.layer4.3.bn2.running_mean',
    'vision.backbone.resnet.layer4.3.bn2.moving_variance':'vision.backbone.resnet.layer4.3.bn2.running_var',
    'vision.backbone.resnet.layer4.3.bn2.gamma':'vision.backbone.resnet.layer4.3.bn2.weight',
    'vision.backbone.resnet.layer4.3.bn2.beta':'vision.backbone.resnet.layer4.3.bn2.bias',
    'vision.backbone.resnet.layer4.4.conv1.weight':'vision.backbone.resnet.layer4.4.conv1.weight',
    'vision.backbone.resnet.layer4.4.bn1.moving_mean':'vision.backbone.resnet.layer4.4.bn1.running_mean',
    'vision.backbone.resnet.layer4.4.bn1.moving_variance':'vision.backbone.resnet.layer4.4.bn1.running_var',
    'vision.backbone.resnet.layer4.4.bn1.gamma':'vision.backbone.resnet.layer4.4.bn1.weight',
    'vision.backbone.resnet.layer4.4.bn1.beta':'vision.backbone.resnet.layer4.4.bn1.bias',
    'vision.backbone.resnet.layer4.4.conv2.weight':'vision.backbone.resnet.layer4.4.conv2.weight',
    'vision.backbone.resnet.layer4.4.bn2.moving_mean':'vision.backbone.resnet.layer4.4.bn2.running_mean',
    'vision.backbone.resnet.layer4.4.bn2.moving_variance':'vision.backbone.resnet.layer4.4.bn2.running_var',
    'vision.backbone.resnet.layer4.4.bn2.gamma':'vision.backbone.resnet.layer4.4.bn2.weight',
    'vision.backbone.resnet.layer4.4.bn2.beta':'vision.backbone.resnet.layer4.4.bn2.bias',
    'vision.backbone.resnet.layer4.5.conv1.weight':'vision.backbone.resnet.layer4.5.conv1.weight',
    'vision.backbone.resnet.layer4.5.bn1.moving_mean':'vision.backbone.resnet.layer4.5.bn1.running_mean',
    'vision.backbone.resnet.layer4.5.bn1.moving_variance':'vision.backbone.resnet.layer4.5.bn1.running_var',
    'vision.backbone.resnet.layer4.5.bn1.gamma':'vision.backbone.resnet.layer4.5.bn1.weight',
    'vision.backbone.resnet.layer4.5.bn1.beta':'vision.backbone.resnet.layer4.5.bn1.bias',
    'vision.backbone.resnet.layer4.5.conv2.weight':'vision.backbone.resnet.layer4.5.conv2.weight',
    'vision.backbone.resnet.layer4.5.bn2.moving_mean':'vision.backbone.resnet.layer4.5.bn2.running_mean',
    'vision.backbone.resnet.layer4.5.bn2.moving_variance':'vision.backbone.resnet.layer4.5.bn2.running_var',
    'vision.backbone.resnet.layer4.5.bn2.gamma':'vision.backbone.resnet.layer4.5.bn2.weight',
    'vision.backbone.resnet.layer4.5.bn2.beta':'vision.backbone.resnet.layer4.5.bn2.bias',
    'vision.backbone.resnet.layer5.0.conv1.weight':'vision.backbone.resnet.layer5.0.conv1.weight',
    'vision.backbone.resnet.layer5.0.bn1.moving_mean':'vision.backbone.resnet.layer5.0.bn1.running_mean',
    'vision.backbone.resnet.layer5.0.bn1.moving_variance':'vision.backbone.resnet.layer5.0.bn1.running_var',
    'vision.backbone.resnet.layer5.0.bn1.gamma':'vision.backbone.resnet.layer5.0.bn1.weight',
    'vision.backbone.resnet.layer5.0.bn1.beta':'vision.backbone.resnet.layer5.0.bn1.bias',
    'vision.backbone.resnet.layer5.0.conv2.weight':'vision.backbone.resnet.layer5.0.conv2.weight',
    'vision.backbone.resnet.layer5.0.bn2.moving_mean':'vision.backbone.resnet.layer5.0.bn2.running_mean',
    'vision.backbone.resnet.layer5.0.bn2.moving_variance':'vision.backbone.resnet.layer5.0.bn2.running_var',
    'vision.backbone.resnet.layer5.0.bn2.gamma':'vision.backbone.resnet.layer5.0.bn2.weight',
    'vision.backbone.resnet.layer5.0.bn2.beta':'vision.backbone.resnet.layer5.0.bn2.bias',
    'vision.backbone.resnet.layer5.0.downsample.0.weight':'vision.backbone.resnet.layer5.0.downsample.0.weight',
    'vision.backbone.resnet.layer5.0.downsample.1.moving_mean':'vision.backbone.resnet.layer5.0.downsample.1.running_mean',
    'vision.backbone.resnet.layer5.0.downsample.1.moving_variance':'vision.backbone.resnet.layer5.0.downsample.1.running_var',
    'vision.backbone.resnet.layer5.0.downsample.1.gamma':'vision.backbone.resnet.layer5.0.downsample.1.weight',
    'vision.backbone.resnet.layer5.0.downsample.1.beta':'vision.backbone.resnet.layer5.0.downsample.1.bias',
    'vision.backbone.resnet.layer5.1.conv1.weight':'vision.backbone.resnet.layer5.1.conv1.weight',
    'vision.backbone.resnet.layer5.1.bn1.moving_mean':'vision.backbone.resnet.layer5.1.bn1.running_mean',
    'vision.backbone.resnet.layer5.1.bn1.moving_variance':'vision.backbone.resnet.layer5.1.bn1.running_var',
    'vision.backbone.resnet.layer5.1.bn1.gamma':'vision.backbone.resnet.layer5.1.bn1.weight',
    'vision.backbone.resnet.layer5.1.bn1.beta':'vision.backbone.resnet.layer5.1.bn1.bias',
    'vision.backbone.resnet.layer5.1.conv2.weight':'vision.backbone.resnet.layer5.1.conv2.weight',
    'vision.backbone.resnet.layer5.1.bn2.moving_mean':'vision.backbone.resnet.layer5.1.bn2.running_mean',
    'vision.backbone.resnet.layer5.1.bn2.moving_variance':'vision.backbone.resnet.layer5.1.bn2.running_var',
    'vision.backbone.resnet.layer5.1.bn2.gamma':'vision.backbone.resnet.layer5.1.bn2.weight',
    'vision.backbone.resnet.layer5.1.bn2.beta':'vision.backbone.resnet.layer5.1.bn2.bias',
    'vision.backbone.resnet.layer5.2.conv1.weight':'vision.backbone.resnet.layer5.2.conv1.weight',
    'vision.backbone.resnet.layer5.2.bn1.moving_mean':'vision.backbone.resnet.layer5.2.bn1.running_mean',
    'vision.backbone.resnet.layer5.2.bn1.moving_variance':'vision.backbone.resnet.layer5.2.bn1.running_var',
    'vision.backbone.resnet.layer5.2.bn1.gamma':'vision.backbone.resnet.layer5.2.bn1.weight',
    'vision.backbone.resnet.layer5.2.bn1.beta':'vision.backbone.resnet.layer5.2.bn1.bias',
    'vision.backbone.resnet.layer5.2.conv2.weight':'vision.backbone.resnet.layer5.2.conv2.weight',
    'vision.backbone.resnet.layer5.2.bn2.moving_mean':'vision.backbone.resnet.layer5.2.bn2.running_mean',
    'vision.backbone.resnet.layer5.2.bn2.moving_variance':'vision.backbone.resnet.layer5.2.bn2.running_var',
    'vision.backbone.resnet.layer5.2.bn2.gamma':'vision.backbone.resnet.layer5.2.bn2.weight',
    'vision.backbone.resnet.layer5.2.bn2.beta':'vision.backbone.resnet.layer5.2.bn2.bias',
    'vision.backbone.transformer.blocks.0.layernorm1.gamma':'vision.backbone.transformer.layers.0.norm1.weight',
    'vision.backbone.transformer.blocks.0.layernorm1.beta':'vision.backbone.transformer.layers.0.norm1.bias',
    'vision.backbone.transformer.blocks.0.layernorm2.gamma':'vision.backbone.transformer.layers.0.norm2.weight',
    'vision.backbone.transformer.blocks.0.layernorm2.beta':'vision.backbone.transformer.layers.0.norm2.bias',
    'vision.backbone.transformer.blocks.0.attention.projection.bias':'vision.backbone.transformer.layers.0.self_attn.out_proj.bias',


    'vision.backbone.transformer.blocks.0.attention.dense1.weight':'vision.backbone.transformer.layers.0.self_attn.in_proj_weight.q',
    'vision.backbone.transformer.blocks.0.attention.dense1.bias':'vision.backbone.transformer.layers.0.self_attn.in_proj_bias.q',
    'vision.backbone.transformer.blocks.0.attention.dense2.weight':'vision.backbone.transformer.layers.0.self_attn.in_proj_weight.k',
    'vision.backbone.transformer.blocks.0.attention.dense2.bias':'vision.backbone.transformer.layers.0.self_attn.in_proj_bias.k',
    'vision.backbone.transformer.blocks.0.attention.dense3.weight':'vision.backbone.transformer.layers.0.self_attn.in_proj_weight.v',
    'vision.backbone.transformer.blocks.0.attention.dense3.bias':'vision.backbone.transformer.layers.0.self_attn.in_proj_bias.v',



    'vision.backbone.transformer.blocks.0.output.mapping.bias':'vision.backbone.transformer.layers.0.linear1.bias',
    'vision.backbone.transformer.blocks.0.output.projection.bias':'vision.backbone.transformer.layers.0.linear2.bias',
    'vision.backbone.transformer.blocks.1.layernorm1.gamma':'vision.backbone.transformer.layers.1.norm1.weight',
    'vision.backbone.transformer.blocks.1.layernorm1.beta':'vision.backbone.transformer.layers.1.norm1.bias',
    'vision.backbone.transformer.blocks.1.layernorm2.gamma':'vision.backbone.transformer.layers.1.norm2.weight',
    'vision.backbone.transformer.blocks.1.layernorm2.beta':'vision.backbone.transformer.layers.1.norm2.bias',
    'vision.backbone.transformer.blocks.1.attention.projection.bias':'vision.backbone.transformer.layers.1.self_attn.out_proj.bias',
    'vision.backbone.transformer.blocks.1.attention.dense1.weight':'vision.backbone.transformer.layers.1.self_attn.in_proj_weight.q',
    'vision.backbone.transformer.blocks.1.attention.dense1.bias':'vision.backbone.transformer.layers.1.self_attn.in_proj_bias.q',
    'vision.backbone.transformer.blocks.1.attention.dense2.weight':'vision.backbone.transformer.layers.1.self_attn.in_proj_weight.k',
    'vision.backbone.transformer.blocks.1.attention.dense2.bias':'vision.backbone.transformer.layers.1.self_attn.in_proj_bias.k',
    'vision.backbone.transformer.blocks.1.attention.dense3.weight':'vision.backbone.transformer.layers.1.self_attn.in_proj_weight.v',
    'vision.backbone.transformer.blocks.1.attention.dense3.bias':'vision.backbone.transformer.layers.1.self_attn.in_proj_bias.v',
    'vision.backbone.transformer.blocks.1.output.mapping.bias':'vision.backbone.transformer.layers.1.linear1.bias',
    'vision.backbone.transformer.blocks.1.output.projection.bias':'vision.backbone.transformer.layers.1.linear2.bias',
    'vision.backbone.transformer.blocks.2.layernorm1.gamma':'vision.backbone.transformer.layers.2.norm1.weight',
    'vision.backbone.transformer.blocks.2.layernorm1.beta':'vision.backbone.transformer.layers.2.norm1.bias',
    'vision.backbone.transformer.blocks.2.layernorm2.gamma':'vision.backbone.transformer.layers.2.norm2.weight',
    'vision.backbone.transformer.blocks.2.layernorm2.beta':'vision.backbone.transformer.layers.2.norm2.bias',
    'vision.backbone.transformer.blocks.2.attention.projection.bias':'vision.backbone.transformer.layers.2.self_attn.out_proj.bias',
    'vision.backbone.transformer.blocks.2.attention.dense1.weight':'vision.backbone.transformer.layers.2.self_attn.in_proj_weight.q',
    'vision.backbone.transformer.blocks.2.attention.dense1.bias':'vision.backbone.transformer.layers.2.self_attn.in_proj_bias.q',
    'vision.backbone.transformer.blocks.2.attention.dense2.weight':'vision.backbone.transformer.layers.2.self_attn.in_proj_weight.k',
    'vision.backbone.transformer.blocks.2.attention.dense2.bias':'vision.backbone.transformer.layers.2.self_attn.in_proj_bias.k',
    'vision.backbone.transformer.blocks.2.attention.dense3.weight':'vision.backbone.transformer.layers.2.self_attn.in_proj_weight.v',
    'vision.backbone.transformer.blocks.2.attention.dense3.bias':'vision.backbone.transformer.layers.2.self_attn.in_proj_bias.v',
    'vision.backbone.transformer.blocks.2.output.mapping.bias':'vision.backbone.transformer.layers.2.linear1.bias',
    'vision.backbone.transformer.blocks.2.output.projection.bias':'vision.backbone.transformer.layers.2.linear2.bias',
    'vision.attention.k_encoder1.0.weight':'vision.attention.k_encoder.0.0.weight',
    'vision.attention.k_encoder1.0.bias':'vision.attention.k_encoder.0.0.bias',
    'vision.attention.k_encoder1.1.moving_mean':'vision.attention.k_encoder.0.1.running_mean',
    'vision.attention.k_encoder1.1.moving_variance':'vision.attention.k_encoder.0.1.running_var',
    'vision.attention.k_encoder1.1.gamma':'vision.attention.k_encoder.0.1.weight',
    'vision.attention.k_encoder1.1.beta':'vision.attention.k_encoder.0.1.bias',
    'vision.attention.k_encoder2.0.weight':'vision.attention.k_encoder.1.0.weight',
    'vision.attention.k_encoder2.0.bias':'vision.attention.k_encoder.1.0.bias',
    'vision.attention.k_encoder2.1.moving_mean':'vision.attention.k_encoder.1.1.running_mean',
    'vision.attention.k_encoder2.1.moving_variance':'vision.attention.k_encoder.1.1.running_var',
    'vision.attention.k_encoder2.1.gamma':'vision.attention.k_encoder.1.1.weight',
    'vision.attention.k_encoder2.1.beta':'vision.attention.k_encoder.1.1.bias',
    'vision.attention.k_encoder3.0.weight':'vision.attention.k_encoder.2.0.weight',
    'vision.attention.k_encoder3.0.bias':'vision.attention.k_encoder.2.0.bias',
    'vision.attention.k_encoder3.1.moving_mean':'vision.attention.k_encoder.2.1.running_mean',
    'vision.attention.k_encoder3.1.moving_variance':'vision.attention.k_encoder.2.1.running_var',
    'vision.attention.k_encoder3.1.gamma':'vision.attention.k_encoder.2.1.weight',
    'vision.attention.k_encoder3.1.beta':'vision.attention.k_encoder.2.1.bias',
    'vision.attention.k_encoder4.0.weight':'vision.attention.k_encoder.3.0.weight',
    'vision.attention.k_encoder4.0.bias':'vision.attention.k_encoder.3.0.bias',
    'vision.attention.k_encoder4.1.moving_mean':'vision.attention.k_encoder.3.1.running_mean',
    'vision.attention.k_encoder4.1.moving_variance':'vision.attention.k_encoder.3.1.running_var',
    'vision.attention.k_encoder4.1.gamma':'vision.attention.k_encoder.3.1.weight',
    'vision.attention.k_encoder4.1.beta':'vision.attention.k_encoder.3.1.bias',
    'vision.attention.k_decoder1.1.weight':'vision.attention.k_decoder.0.1.weight',
    'vision.attention.k_decoder1.1.bias':'vision.attention.k_decoder.0.1.bias',
    'vision.attention.k_decoder1.2.moving_mean':'vision.attention.k_decoder.0.2.running_mean',
    'vision.attention.k_decoder1.2.moving_variance':'vision.attention.k_decoder.0.2.running_var',
    'vision.attention.k_decoder1.2.gamma':'vision.attention.k_decoder.0.2.weight',
    'vision.attention.k_decoder1.2.beta':'vision.attention.k_decoder.0.2.bias',
    'vision.attention.k_decoder2.1.weight':'vision.attention.k_decoder.1.1.weight',
    'vision.attention.k_decoder2.1.bias':'vision.attention.k_decoder.1.1.bias',
    'vision.attention.k_decoder2.2.moving_mean':'vision.attention.k_decoder.1.2.running_mean',
    'vision.attention.k_decoder2.2.moving_variance':'vision.attention.k_decoder.1.2.running_var',
    'vision.attention.k_decoder2.2.gamma':'vision.attention.k_decoder.1.2.weight',
    'vision.attention.k_decoder2.2.beta':'vision.attention.k_decoder.1.2.bias',
    'vision.attention.k_decoder3.1.weight':'vision.attention.k_decoder.2.1.weight',
    'vision.attention.k_decoder3.1.bias':'vision.attention.k_decoder.2.1.bias',
    'vision.attention.k_decoder3.2.moving_mean':'vision.attention.k_decoder.2.2.running_mean',
    'vision.attention.k_decoder3.2.moving_variance':'vision.attention.k_decoder.2.2.running_var',
    'vision.attention.k_decoder3.2.gamma':'vision.attention.k_decoder.2.2.weight',
    'vision.attention.k_decoder3.2.beta':'vision.attention.k_decoder.2.2.bias',
    'vision.attention.k_decoder4.1.weight':'vision.attention.k_decoder.3.1.weight',
    'vision.attention.k_decoder4.1.bias':'vision.attention.k_decoder.3.1.bias',
    'vision.attention.k_decoder4.2.moving_mean':'vision.attention.k_decoder.3.2.running_mean',
    'vision.attention.k_decoder4.2.moving_variance':'vision.attention.k_decoder.3.2.running_var',
    'vision.attention.k_decoder4.2.gamma':'vision.attention.k_decoder.3.2.weight',
    'vision.attention.k_decoder4.2.beta':'vision.attention.k_decoder.3.2.bias',
    'vision.attention.project.weight':'vision.attention.project.weight',
    'vision.attention.project.bias':'vision.attention.project.bias',
    'vision.cls.weight':'vision.cls.weight',
    'vision.cls.bias':'vision.cls.bias',
    'vision.backbone.pos_encoder.pe':'vision.backbone.pos_encoder.pe',
    'vision.attention.pos_encoder.pe':'vision.attention.pos_encoder.pe',
    'language.proj.weight':'language.proj.weight',
    'language.model.blocks.0.cross_attention.dense1.weight':'language.model.layers.0.multihead_attn.in_proj_weight.q',
    'language.model.blocks.0.cross_attention.dense1.bias':'language.model.layers.0.multihead_attn.in_proj_bias.q',
    'language.model.blocks.0.cross_attention.dense2.weight':'language.model.layers.0.multihead_attn.in_proj_weight.k',
    'language.model.blocks.0.cross_attention.dense2.bias':'language.model.layers.0.multihead_attn.in_proj_bias.k',
    'language.model.blocks.0.cross_attention.dense3.weight':'language.model.layers.0.multihead_attn.in_proj_weight.v',
    'language.model.blocks.0.cross_attention.dense3.bias':'language.model.layers.0.multihead_attn.in_proj_bias.v',
    'language.model.blocks.0.cross_attention.projection.bias':'language.model.layers.0.multihead_attn.out_proj.bias',
    'language.model.blocks.0.cross_attention_layernorm.gamma':'language.model.layers.0.norm2.weight',
    'language.model.blocks.0.cross_attention_layernorm.beta':'language.model.layers.0.norm2.bias',
    'language.model.blocks.0.layernorm2.gamma':'language.model.layers.0.norm3.weight',
    'language.model.blocks.0.layernorm2.beta':'language.model.layers.0.norm3.bias',
    'language.model.blocks.0.output.projection.bias':'language.model.layers.0.linear2.bias',
    'language.model.blocks.0.output.mapping.bias':'language.model.layers.0.linear1.bias',
    'language.model.blocks.1.cross_attention.dense1.weight':'language.model.layers.1.multihead_attn.in_proj_weight.q',
    'language.model.blocks.1.cross_attention.dense1.bias':'language.model.layers.1.multihead_attn.in_proj_bias.q',
    'language.model.blocks.1.cross_attention.dense2.weight':'language.model.layers.1.multihead_attn.in_proj_weight.k',
    'language.model.blocks.1.cross_attention.dense2.bias':'language.model.layers.1.multihead_attn.in_proj_bias.k',
    'language.model.blocks.1.cross_attention.dense3.weight':'language.model.layers.1.multihead_attn.in_proj_weight.v',
    'language.model.blocks.1.cross_attention.dense3.bias':'language.model.layers.1.multihead_attn.in_proj_bias.v',
    'language.model.blocks.1.cross_attention.projection.bias':'language.model.layers.1.multihead_attn.out_proj.bias',
    'language.model.blocks.1.cross_attention_layernorm.gamma':'language.model.layers.1.norm2.weight',
    'language.model.blocks.1.cross_attention_layernorm.beta':'language.model.layers.1.norm2.bias',
    'language.model.blocks.1.layernorm2.gamma':'language.model.layers.1.norm3.weight',
    'language.model.blocks.1.layernorm2.beta':'language.model.layers.1.norm3.bias',
    'language.model.blocks.1.output.projection.bias':'language.model.layers.1.linear2.bias',
    'language.model.blocks.1.output.mapping.bias':'language.model.layers.1.linear1.bias',
    'language.model.blocks.2.cross_attention.dense1.weight':'language.model.layers.2.multihead_attn.in_proj_weight.q',
    'language.model.blocks.2.cross_attention.dense1.bias':'language.model.layers.2.multihead_attn.in_proj_bias.q',
    'language.model.blocks.2.cross_attention.dense2.weight':'language.model.layers.2.multihead_attn.in_proj_weight.k',
    'language.model.blocks.2.cross_attention.dense2.bias':'language.model.layers.2.multihead_attn.in_proj_bias.k',
    'language.model.blocks.2.cross_attention.dense3.weight':'language.model.layers.2.multihead_attn.in_proj_weight.v',
    'language.model.blocks.2.cross_attention.dense3.bias':'language.model.layers.2.multihead_attn.in_proj_bias.v',
    'language.model.blocks.2.cross_attention.projection.bias':'language.model.layers.2.multihead_attn.out_proj.bias',
    'language.model.blocks.2.cross_attention_layernorm.gamma':'language.model.layers.2.norm2.weight',
    'language.model.blocks.2.cross_attention_layernorm.beta':'language.model.layers.2.norm2.bias',
    'language.model.blocks.2.layernorm2.gamma':'language.model.layers.2.norm3.weight',
    'language.model.blocks.2.layernorm2.beta':'language.model.layers.2.norm3.bias',
    'language.model.blocks.2.output.projection.bias':'language.model.layers.2.linear2.bias',
    'language.model.blocks.2.output.mapping.bias':'language.model.layers.2.linear1.bias',   
    'language.model.blocks.3.cross_attention.dense1.weight':'language.model.layers.3.multihead_attn.in_proj_weight.q',
    'language.model.blocks.3.cross_attention.dense1.bias':'language.model.layers.3.multihead_attn.in_proj_bias.q',
    'language.model.blocks.3.cross_attention.dense2.weight':'language.model.layers.3.multihead_attn.in_proj_weight.k',
    'language.model.blocks.3.cross_attention.dense2.bias':'language.model.layers.3.multihead_attn.in_proj_bias.k',
    'language.model.blocks.3.cross_attention.dense3.weight':'language.model.layers.3.multihead_attn.in_proj_weight.v',
    'language.model.blocks.3.cross_attention.dense3.bias':'language.model.layers.3.multihead_attn.in_proj_bias.v',
    'language.model.blocks.3.cross_attention.projection.bias':'language.model.layers.3.multihead_attn.out_proj.bias',
    'language.model.blocks.3.cross_attention_layernorm.gamma':'language.model.layers.3.norm2.weight',
    'language.model.blocks.3.cross_attention_layernorm.beta':'language.model.layers.3.norm2.bias',
    'language.model.blocks.3.layernorm2.gamma':'language.model.layers.3.norm3.weight',
    'language.model.blocks.3.layernorm2.beta':'language.model.layers.3.norm3.bias',
    'language.model.blocks.3.output.projection.bias':'language.model.layers.3.linear2.bias',
    'language.model.blocks.3.output.mapping.bias':'language.model.layers.3.linear1.bias',
    'language.cls.weight':'language.cls.weight',
    'language.cls.bias':'language.cls.bias',
    'alignment.w_att.weight':'alignment.w_att.weight',
    'alignment.w_att.bias':'alignment.w_att.bias',
    'alignment.cls.weight':'alignment.cls.weight',
    'alignment.cls.bias':'alignment.cls.bias',
    'language.token_encoder.pe':'language.token_encoder.pe',
    'language.pos_encoder.pe':'language.pos_encoder.pe'
    }

    parm1 = {
        'vision.backbone.transformer.blocks.2.output.mapping.weight':'vision.backbone.transformer.layers.2.linear1.weight',
        'vision.backbone.transformer.blocks.2.output.projection.weight':'vision.backbone.transformer.layers.2.linear2.weight',
        'vision.backbone.transformer.blocks.1.output.projection.weight':'vision.backbone.transformer.layers.1.linear2.weight',
        'vision.backbone.transformer.blocks.1.output.mapping.weight':'vision.backbone.transformer.layers.1.linear1.weight',
        'vision.backbone.transformer.blocks.0.output.projection.weight':'vision.backbone.transformer.layers.0.linear2.weight',
        'vision.backbone.transformer.blocks.0.output.mapping.weight':'vision.backbone.transformer.layers.0.linear1.weight',
        'vision.backbone.transformer.blocks.0.attention.projection.weight':'vision.backbone.transformer.layers.0.self_attn.out_proj.weight',
        'vision.backbone.transformer.blocks.1.attention.projection.weight':'vision.backbone.transformer.layers.1.self_attn.out_proj.weight',
        'vision.backbone.transformer.blocks.2.attention.projection.weight':'vision.backbone.transformer.layers.2.self_attn.out_proj.weight',
        'language.model.blocks.0.cross_attention.projection.weight':'language.model.layers.0.multihead_attn.out_proj.weight',
        'language.model.blocks.0.output.mapping.weight':'language.model.layers.0.linear1.weight',
        'language.model.blocks.0.output.projection.weight':'language.model.layers.0.linear2.weight',
        'language.model.blocks.1.cross_attention.projection.weight':'language.model.layers.1.multihead_attn.out_proj.weight',
        'language.model.blocks.1.output.mapping.weight':'language.model.layers.1.linear1.weight',
        'language.model.blocks.1.output.projection.weight':'language.model.layers.1.linear2.weight',
        'language.model.blocks.2.cross_attention.projection.weight':'language.model.layers.2.multihead_attn.out_proj.weight',
        'language.model.blocks.2.output.mapping.weight':'language.model.layers.2.linear1.weight',
        'language.model.blocks.2.output.projection.weight':'language.model.layers.2.linear2.weight',
        'language.model.blocks.3.cross_attention.projection.weight':'language.model.layers.3.multihead_attn.out_proj.weight',
        'language.model.blocks.3.output.mapping.weight':'language.model.layers.3.linear1.weight',
        'language.model.blocks.3.output.projection.weight':'language.model.layers.3.linear2.weight',
    }


    new_params_list = []
    for ms_param in ms_params.keys():
        if('vision.attention.k_encoder2.0.weight' in ms_param):
            print("1")
            print(pt_params[param[ms_param]].shape)
            print(ms_params[ms_param].shape)
        if ms_param in param and pt_params[param[ms_param]].shape == ms_params[ms_param].shape:


            ms_value = pt_params[param[ms_param]]
            new_params_list.append({"name": ms_param, "data": ms.Tensor(ms_value)})
        elif ms_param in parm1 :#Mindspore和Pytorch
            # if('language.model.blocks.0.' in ms_param):
            #     print("111")
            ms_value = pt_params[parm1[ms_param]]
            ms_value = ms_value.transpose(1,0)
            if ms_value.shape == ms_params[ms_param].shape:
                new_params_list.append({"name": ms_param, "data": ms.Tensor(ms_value)})
            else:
                print(ms_param, "not match in pt_params")
        else:
            print(ms_param, "not match in pt_params")



    # 保存成MindSpore的checkpoint
    ms.save_checkpoint(new_params_list, ckpt_path)

config = Config("/home/data/zhangyh22/ABINet/ABINet_Mindspore/abinet_gpu/abinet_ascend/utils/configs/train_abinet.yaml")
pth_path = "/home/data/zhangyh22/ABINet/ABINet_Mindspore/workdir/train-abinet/best-train-abinet.pth"
pt_param = pytorch_params_convert_weight(pth_path)
print("="*20)
ms_abinet = ms_ABINetIterModel(config)
print("1111111111111")
ms_param = mindspore_params(ms_abinet)

split_torch_attention(pt_param)
print("split")
ckpt_path = "./ms_abinet_from_py_ascend.ckpt"
param_convert(ms_param, pt_param, ckpt_path)
print("done")




# def get_converted_ckpt(mapped_params, weight_dict):
#     """
#     Print the keys of the loaded checkpoint

#     Args:
#         mapped_params(dict): The loaded checkpoint. The key is parameter name and value is the numpy array.
#         weight_dict(dict): The loaded pytorch checkpoint.

#     Returns:
#         None
#     """
#     new_ckpt_list = []
#     print(weight_dict.keys())
#     # Currently, the ms_extend_param the torch_extend_param is the full parameters.
#     for src, tgt in mapped_params:
#         value = weight_dict[tgt].numpy()
#         if '.dense.weight' in tgt:
#             value = np.transpose(value, [1, 0])

#         print(f"Mapping table Mindspore:{src:<30} \t Torch:{tgt:<30} with shape {value.shape}")

#         new_ckpt_list.append({"data": Tensor(value), "name": src})
#     return new_ckpt_list


# # def split_torch_attention(state):
# #     s = list(state.keys())
# #     for name in s:
# #         if name.endswith('attn.c_attn.weight') or name.endswith('attn.c_attn.bias'):
# #             value = state.pop(name)
# #             q, k, v = np.split(value.numpy(), 3, -1)
# #             state[name + '.q'] = torch.tensor(q, dtype=value.dtype)
# #             state[name + '.k'] = torch.tensor(k)
# #             state[name + '.v'] = torch.tensor(v)


# if __name__ == '__main__':
#     parser = argparse.ArgumentParser(description="BERT convert script")
#     parser.add_argument('--layers',
#                         type=int,
#                         default=1,
#                         help="The number of layers of the model to be converted.")
#     parser.add_argument("--torch_path",
#                         type=str,
#                         default=None,
#                         required=True,
#                         help="The torch checkpoint path.")
#     parser.add_argument("--mindspore_path",
#                         type=str,
#                         required=True,
#                         default="The output mindspore checkpoint path.",
#                         help="Use device nums, default is 128.")

#     opt = parser.parse_args()
#     state_dict = torch.load(opt.torch_path, map_location='cpu')
#     print_dict(state_dict)

#     ms_name = [
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense1.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense1.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense2.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense2.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense3.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.dense3.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.projection.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.attention.projection.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.layernorm2.gamma",
#         "bert.bert.bert_encoder.encoder.blocks.{}.layernorm2.beta",
#         "bert.bert.bert_encoder.encoder.blocks.{}.output.mapping.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.output.mapping.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.output.projection.weight",
#         "bert.bert.bert_encoder.encoder.blocks.{}.output.projection.bias",
#         "bert.bert.bert_encoder.encoder.blocks.{}.layernorm1.gamma",
#         "bert.bert.bert_encoder.encoder.blocks.{}.layernorm1.beta",
#     ]

#     torch_name = [
#         "bert.encoder.layer.{}.attention.self.query.weight",
#         "bert.encoder.layer.{}.attention.self.query.bias",
#         "bert.encoder.layer.{}.attention.self.key.weight",
#         "bert.encoder.layer.{}.attention.self.key.bias",
#         "bert.encoder.layer.{}.attention.self.value.weight",
#         "bert.encoder.layer.{}.attention.self.value.bias",
#         "bert.encoder.layer.{}.attention.output.dense.weight",
#         "bert.encoder.layer.{}.attention.output.dense.bias",
#         "bert.encoder.layer.{}.attention.output.LayerNorm.gamma",
#         "bert.encoder.layer.{}.attention.output.LayerNorm.beta",
#         "bert.encoder.layer.{}.intermediate.dense.weight",
#         "bert.encoder.layer.{}.intermediate.dense.bias",
#         "bert.encoder.layer.{}.output.dense.weight",
#         "bert.encoder.layer.{}.output.dense.bias",
#         "bert.encoder.layer.{}.output.LayerNorm.gamma",
#         "bert.encoder.layer.{}.output.LayerNorm.beta",
#     ]

#     addition_mindspore = [
#         "bert.bert.word_embedding.embedding_table",
#         "bert.bert.embedding_postprocessor.full_position_embedding.embedding_table",
#         "bert.bert.embedding_postprocessor.token_type_embedding.embedding_table",
#         "bert.bert.embedding_postprocessor.layernorm.gamma",
#         "bert.bert.embedding_postprocessor.layernorm.beta",
#         "bert.bert.dense.weight",
#         "bert.bert.dense.bias",
#     ]

#     addition_torch = [
#         "bert.embeddings.word_embeddings.weight",
#         "bert.embeddings.position_embeddings.weight",
#         "bert.embeddings.token_type_embeddings.weight",
#         "bert.embeddings.LayerNorm.gamma",
#         "bert.embeddings.LayerNorm.beta",
#         "bert.pooler.dense.weight",
#         "bert.pooler.dense.bias",
#     ]

#     mapped_param = generate_params_dict(total_layers=opt.layers,
#                                         mindspore_params_per_layer=ms_name,
#                                         torch_params_per_layer=torch_name,
#                                         mindspore_additional_params=addition_mindspore,
#                                         torch_additional_params=addition_torch)
#     split_torch_attention(state_dict)
#     new_ckpt = get_converted_ckpt(mapped_param, state_dict)
#     save_checkpoint(new_ckpt, opt.mindspore_path)
#     print(f"Convert finished, the output is saved to {opt.mindspore_path}")



